"""
轻量级翻译服务器
使用小型但可靠的翻译模型
"""

from flask import Flask, request, jsonify
from flask_cors import CORS
import sys
import subprocess

app = Flask(__name__)
CORS(app)

# 全局变量
model = None
tokenizer = None
model_name = None

def install_packages():
    """安装必要的包"""
    try:
        print("检查并安装必要的包...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers", "torch", "sentencepiece", "-q"])
        print("✓ 包安装完成")
        return True
    except Exception as e:
        print(f"✗ 包安装失败: {e}")
        return False

def load_model():
    """
    加载轻量级翻译模型
    尝试多个小型模型
    """
    global model, tokenizer, model_name
    
    # 按大小从小到大尝试多个模型
    models_to_try = [
        "Helsinki-NLP/opus-mt-en-zh",  # 约312MB
    ]
    
    for model_name_try in models_to_try:
        try:
            print(f"\n尝试加载模型: {model_name_try}")
            from transformers import MarianMTModel, MarianTokenizer
            
            tokenizer = MarianTokenizer.from_pretrained(model_name_try)
            model = MarianMTModel.from_pretrained(model_name_try)
            model_name = model_name_try
            
            print(f"✓ 成功加载模型: {model_name_try}")
            return True
            
        except Exception as e:
            print(f"✗ 模型 {model_name_try} 加载失败: {e}")
            continue
    
    return False

def translate_text(text):
    """执行翻译"""
    if model is None or tokenizer is None:
        return "模型未加载"
    
    try:
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
        translated = model.generate(**inputs, max_length=512, num_beams=4, early_stopping=True)
        translation = tokenizer.decode(translated[0], skip_special_tokens=True)
        return translation
    except Exception as e:
        return f"翻译出错: {str(e)}"

@app.route('/translate', methods=['POST'])
def translate():
    """翻译API端点"""
    try:
        data = request.json
        text = data.get('text', '')
        
        if not text:
            return jsonify({'error': '文本为空'}), 400
        
        translation = translate_text(text)
        return jsonify({'translation': translation})
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/health', methods=['GET'])
def health():
    """健康检查端点"""
    return jsonify({
        'status': 'ready' if model is not None else 'loading',
        'model': model_name or 'none',
        'model_loaded': model is not None
    })

if __name__ == '__main__':
    print("=" * 60)
    print("轻量级翻译服务器")
    print("=" * 60)
    
    print("\n步骤1: 安装依赖包...")
    if not install_packages():
        print("✗ 依赖安装失败")
        sys.exit(1)
    
    print("\n步骤2: 加载翻译模型...")
    if not load_model():
        print("✗ 所有模型加载失败")
        sys.exit(1)
    
    print("\n✓ 服务器准备就绪")
    print(f"✓ 使用模型: {model_name}")
    print("\n服务器地址: http://localhost:5000")
    print("翻译接口: POST http://localhost:5000/translate")
    print("健康检查: GET http://localhost:5000/health")
    print("\n按 Ctrl+C 停止服务器")
    print("=" * 60)
    
    app.run(host='0.0.0.0', port=5000, debug=False)